"""根据 configs/default.yaml 生成指定配置的 VQModel config.yaml
"""

import yaml
from omegaconf import OmegaConf
from omegaconf.dictconfig import DictConfig
from pathlib import Path

from argparse import ArgumentParser


def prepare_args():
    parser=ArgumentParser()
    parser.add_argument(
        "--save_dir",
        type=str,
    )
    
    parser.add_argument(
        "--codebook_dim",
        type=int,
    )
    parser.add_argument(
        "--codebook_size",
        type=int,
    )
    parser.add_argument(
        "--vae_latent_dim",
        type=int,
    )
    parser.add_argument(
        "--base_ch",
        type=int,
        default=128,
    )
    parser.add_argument(
        "--ch_mult",
        type=str,
        help="format like '1,2,4' , comma seperated"
    )
    parser.add_argument(
        "--attn_resolutions",
        type=str,
        help="usually the lowest resolution or none ; format like '16,32' , comma seperated"
    )
    parser.add_argument(
        "--disc_start",
        type=int,
        default=0,
    )
    parser.add_argument(
        "--disc_num_layers",
        type=int,
        default=3,
    )
    parser.add_argument(
        "--disc_weight",
        type=float,
        default=0.75,
    )
    parser.add_argument(
        "--bs",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--dl_workers",
        type=int,
        default=4,
    )
    parser.add_argument(
        "--data_root",
        type=str,
        default="data/OpenImages",
    )
    parser.add_argument(
        "--train_images_list_file",
        type=str,
        default="data/OpenImages/custom_train_1743042.txt",
    )
    parser.add_argument(
        "--val_images_list_file",
        type=str,
        default="data/OpenImages/custom_val_41620.txt",
    )
    parser.add_argument(
        "--size",
        type=int,
        default=0,
        help="intermediate img resolution ; you may first resize your image to 'size' (shorter edge, remain aspect ratio) , then crop it ; or if no crop, just resize your image to (size,size) "
    )
    parser.add_argument(
        "--crop_size",
        type=int,
        default=256,
        help="crop your image to (crop_size, crop_size) ; if 0 , no crop , and to make sure 'size'>0 (image will be resize to (size,size) )",
    )
    parser.add_argument(
        "--ckpt_path",
        type=str,
        default=None
    )
    parser.add_argument(
        "--no_attn",
        action="store_true",
    )
    args=parser.parse_args()
    
    # tolist
    args.ch_mult = [int(i) for i in (args.ch_mult).rstrip(',').split(',')]
    if args.attn_resolutions=='':
        args.attn_resolutions=[]
    else:
        args.attn_resolutions = [int(i) for i in (args.attn_resolutions).rstrip(',').split(',')]
    
    return args

if __name__ == "__main__":
    args=prepare_args()
    base_cfg=OmegaConf.load("configs/default.yaml")

    final_res=None
    if args.size==0:
        assert args.crop_size>0
        final_res=args.crop_size
    else:
        assert args.crop_size<=args.size
        if args.crop_size==0:
            final_res=args.size
        else:
            final_res=args.crop_size
            
    base_cfg.res = final_res
        
    # base_cfg.res = args.resolution
    base_cfg.model.params.embed_dim = args.codebook_dim
    base_cfg.model.params.n_embed = args.codebook_size
    if args.ckpt_path is not None:
        base_cfg.model.params.ckpt_path = args.ckpt_path
    
    ddconfig = base_cfg.model.params.ddconfig
    ddconfig.resolution=final_res
    ddconfig.z_channels=args.vae_latent_dim
    ddconfig.ch=args.base_ch
    ddconfig.ch_mult=args.ch_mult
    ddconfig.attn_resolutions=args.attn_resolutions
    if args.no_attn:
        ddconfig.attn_type='none'
    
    lossconfig = base_cfg.model.params.lossconfig.params
    lossconfig.disc_start = args.disc_start
    lossconfig.disc_num_layers = args.disc_num_layers
    lossconfig.disc_weight = args.disc_weight
    
    base_cfg.data.params.batch_size=args.bs
    base_cfg.data.params.num_workers=args.dl_workers
    train_ds, val_ds = base_cfg.data.params.train.params, base_cfg.data.params.validation.params
    
    train_ds.data_root=args.data_root
    train_ds.images_list_file=args.train_images_list_file
    train_ds.size=args.size
    train_ds.crop_size=args.crop_size
    
    val_ds.data_root=args.data_root
    val_ds.images_list_file=args.val_images_list_file
    val_ds.size=args.size
    val_ds.crop_size=args.crop_size
    
    # if args.resize_resolution:
    #     train_ds.size=args.resize_resolution
    #     val_ds.size=args.resize_resolution
    #     train_ds.crop_size=args.resolution
    #     val_ds.crop_size=args.resolution
    
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    OmegaConf.save(base_cfg, str(save_dir / 'base.yaml'), )